package com.ai.springai.controller;

import com.ai.springai.entity.Message;
import com.ai.springai.util.TransformerUtil;
import groovy.util.logging.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import static org.springframework.ai.chat.messages.MessageType.*;

@Slf4j
@RestController
@RequestMapping("/ai/chat")
public class ChatController {

    private static final Logger log = LoggerFactory.getLogger(ChatController.class);
    private final OllamaChatClient chatClient;


    @Autowired
    public ChatController(OllamaChatClient chatClient) {
        this.chatClient = chatClient;
    }

    @PostMapping("/chat")
    public String chat(@RequestBody List<Message> message) throws IOException{
//        FileInputStream inputStream = new FileInputStream("F:\\sjj.txt");
        List<Document> documentList = List.of(new Document(TransformerUtil.pdfToStr(new File("F:\\test.pdf"))));
        List<org.springframework.ai.chat.messages.Message> messageList = new ArrayList<>(getSystemMessages(documentList));
        message.forEach(msg->{
            if (USER.equals(msg.getMessageType())) {
                messageList.add(new UserMessage(msg.getContent()));
            }else{
                messageList.add(new AssistantMessage(msg.getContent()));
            }
        });
        Prompt prompt = new Prompt(messageList);
        return chatClient.call(prompt).getResult().getOutput().getContent();
    }

    @GetMapping(value = "/stream", produces = "text/event-stream;charset=utf-8")
    public Flux<String> stream(String message) {
        return chatClient.stream(message);
    }

    private List<SystemMessage> getSystemMessages(List<Document> documentList) {
        List<SystemMessage> messageList = new ArrayList<>();
        documentList.forEach(document ->
            messageList.add(new SystemMessage(document.getContent()))
        );
        return messageList;
    }

}
